import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.colors as mcolors
from scipy.interpolate import griddata

import numpy as np
import matplotlib.pyplot as plt
from pyswarm import pso
from scipy.optimize import minimize
from scipy.optimize import Bounds
import matplotlib as mpl
import matplotlib.font_manager as fm

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

g = 9.8
missile_v = 300
cloud_sink = 3
cloud_r = 10
cloud_t = 20

fake_target = np.array([0, 0, 0])
real_target = np.array([0, 200, 0])
target_r = 7
target_h = 10

M1_start = np.array([20000, 0, 2000])
FY1_start = np.array([17800, 0, 1800])

missile_dir = fake_target - M1_start
missile_dir_norm = np.linalg.norm(missile_dir)
missile_unit_dir = missile_dir / missile_dir_norm
missile_time = missile_dir_norm / missile_v

def missile_position(t):
    return M1_start + missile_unit_dir * missile_v * t

def cloud_position(alpha, v, t_drop, t_burst, tau):
    drop_pos = FY1_start + np.array([v * np.cos(alpha) * t_drop,
                                     v * np.sin(alpha) * t_drop,
                                     0])

    burst_pos = drop_pos + np.array([v * np.cos(alpha) * t_burst,
                                     v * np.sin(alpha) * t_burst,
                                     -0.5 * g * t_burst ** 2])

    return np.array([burst_pos[0],
                     burst_pos[1],
                     max(0, burst_pos[2] - cloud_sink * tau)])

def distance_to_line(point, line_start, line_end):
    line_vec = line_end - line_start
    point_vec = point - line_start
    line_len = np.linalg.norm(line_vec)

    if line_len < 1e-9:  # 避免除以零
        return np.linalg.norm(point - line_start)

    line_unitvec = line_vec / line_len
    point_vec_scaled = point_vec / line_len
    t = np.dot(line_unitvec, point_vec_scaled)

    if t < 0:
        t = 0
    elif t > 1:
        t = 1

    nearest = line_start + t * line_vec
    return np.linalg.norm(nearest - point)

def sight_blocked(missile_pos, cloud_pos, target_center, target_r, target_h):
    num_points = 8
    target_points = []

    for i in range(num_points):
        angle = 2 * np.pi * i / num_points
        x = target_center[0] + target_r * np.cos(angle)
        y = target_center[1] + target_r * np.sin(angle)
        z = target_center[2]
        target_points.append(np.array([x, y, z]))

    for i in range(num_points):
        angle = 2 * np.pi * i / num_points
        x = target_center[0] + target_r * np.cos(angle)
        y = target_center[1] + target_r * np.sin(angle)
        z = target_center[2] + target_h
        target_points.append(np.array([x, y, z]))

    max_dist = 0
    for target_point in target_points:
        dist = distance_to_line(cloud_pos, missile_pos, target_point)
        if dist > max_dist:
            max_dist = dist
            if max_dist > cloud_r:
                return False
    return max_dist <= cloud_r


def effective_shielding_time(alpha, v, t_drop, t_burst):
    t_bur_total = t_drop + t_burst
    t_cloud_end = t_bur_total + cloud_t
    if t_cloud_end > missile_time:
        t_cloud_end = missile_time
    time_step = 0.5
    effective_time = 0
    for t in np.arange(0, min(cloud_t, missile_time - t_bur_total), time_step):
        t_current = t_bur_total + t
        missile_pos = missile_position(t_current)
        cloud_pos = cloud_position(alpha, v, t_drop, t_burst, t)
        if sight_blocked(missile_pos, cloud_pos, real_target, target_r, target_h):
            effective_time += time_step
    return effective_time

def objective_function(x):
    alpha, v, t_drop, t_burst = x
    if (v < 70 or v > 140 or
            t_drop < 0 or t_burst < 0 or
            t_drop + t_burst > missile_time - cloud_t):
        return 1e6
    shielding_time = effective_shielding_time(alpha, v, t_drop, t_burst)
    if shielding_time == 0:
        missile_pos_at_burst = missile_position(t_drop + t_burst)
        cloud_pos_at_burst = cloud_position(alpha, v, t_drop, t_burst, 0)

        missile_to_target_dist = np.linalg.norm(real_target - missile_pos_at_burst)
        cloud_to_target_dist = np.linalg.norm(real_target - cloud_pos_at_burst)

        penalty = 1000 * (1 / (missile_to_target_dist + cloud_to_target_dist + 1))
        return penalty
    return -shielding_time

lb = [0, 70, 0, 0]
ub = [2 * np.pi, 140, missile_time - cloud_t, cloud_t]
np.random.seed(42)

print("开始PSO优化...")
xopt, fopt = pso(objective_function, lb, ub, swarmsize=50, maxiter=100, minstep=1e-4)
print("PSO优化完成，开始局部搜索...")
bounds = Bounds(lb, ub)
result = minimize(objective_function, xopt, method='SLSQP', bounds=bounds)
xopt_refined = result.x
fopt_refined = result.fun

alpha_opt, v_opt, t_drop_opt, t_burst_opt = xopt_refined
def precise_effective_shielding_time(alpha, v, t_drop, t_burst):
    t_burst_total = t_drop + t_burst
    t_cloud_end = t_burst_total + cloud_t
    if t_cloud_end > missile_time:
        t_cloud_end = missile_time
    time_step = 0.1
    effective_time = 0
    for t in np.arange(0, min(cloud_t, missile_time - t_burst_total), time_step):
        t_current = t_burst_total + t
        missile_pos = missile_position(t_current)
        cloud_pos = cloud_position(alpha, v, t_drop, t_burst, t)
        if sight_blocked(missile_pos, cloud_pos, real_target, target_r, target_h):
            effective_time += time_step
    return effective_time

precise_shielding_time = precise_effective_shielding_time(alpha_opt, v_opt, t_drop_opt, t_burst_opt)
print("PSO优化结果:")
print(f"飞行方向: {np.degrees(xopt[0]):.2f}°")
print(f"飞行速度: {xopt[1]:.2f} m/s")
print(f"投放时间: {xopt[2]:.2f} s")
print(f"起爆时间: {xopt[3]:.2f} s")
print(f"最大有效遮蔽时间: {-fopt:.2f} s")

print("\n局部搜索优化结果:")
print(f"飞行方向: {np.degrees(xopt_refined[0]):.2f}°")
print(f"飞行速度: {xopt_refined[1]:.2f} m/s")
print(f"投放时间: {xopt_refined[2]:.2f} s")
print(f"起爆时间: {xopt_refined[3]:.2f} s")
print(f"最大有效遮蔽时间: {-fopt_refined:.2f} s")

print(f"\n精确计算的有效遮蔽时间: {precise_shielding_time:.2f} s")

drop_point = FY1_start + np.array([v_opt * np.cos(alpha_opt) * t_drop_opt,
                                   v_opt * np.sin(alpha_opt) * t_drop_opt,
                                   0])
burst_point = drop_point + np.array([v_opt * np.cos(alpha_opt) * t_burst_opt,
                                     v_opt * np.sin(alpha_opt) * t_burst_opt,
                                     -0.5 * g * t_burst_opt ** 2])

print(f"\n投放点坐标: ({drop_point[0]:.2f}, {drop_point[1]:.2f}, {drop_point[2]:.2f})")
print(f"起爆点坐标: ({burst_point[0]:.2f}, {burst_point[1]:.2f}, {burst_point[2]:.2f})")



# 使用之前优化得到的最佳参数
alpha_opt, v_opt, t_drop_opt, t_burst_opt = xopt_refined

# 在最佳参数附近采样100个点
n_samples = 100
np.random.seed(42)

# 定义采样范围（最佳值的±10%）
alpha_range = [alpha_opt * 0.9, alpha_opt * 1.1]
v_range = [v_opt * 0.9, v_opt * 1.1]
t_drop_range = [t_drop_opt * 0.9, t_drop_opt * 1.1]
t_burst_range = [t_burst_opt * 0.9, t_burst_opt * 1.1]

# 生成随机采样点
samples = []
for i in range(n_samples):
    alpha_sample = np.random.uniform(alpha_range[0], alpha_range[1])
    v_sample = np.random.uniform(v_range[0], v_range[1])
    t_drop_sample = np.random.uniform(t_drop_range[0], t_drop_range[1])
    t_burst_sample = np.random.uniform(t_burst_range[0], t_burst_range[1])
    samples.append([alpha_sample, v_sample, t_drop_sample, t_burst_sample])

samples = np.array(samples)

# 计算每个采样点的遮蔽时间
shielding_times = []
print("计算100个采样点的遮蔽时间...")
for i, sample in enumerate(samples):
    if i % 10 == 0:
        print(f"计算第 {i + 1}/{n_samples} 个点...")
    shielding_time = precise_effective_shielding_time(sample[0], sample[1], sample[2], sample[3])
    shielding_times.append(shielding_time)

shielding_times = np.array(shielding_times)

print("计算完成！")

# 可视化结果
fig = plt.figure(figsize=(10, 7))

# 1. 2D散点图 - 飞行方向 vs 飞行速度
# ax1 = fig.add_subplot(221)
# sc1 = ax1.scatter(np.degrees(samples[:, 0]), samples[:, 1],
#                   c=shielding_times, cmap='viridis', s=50, alpha=0.8)
# ax1.set_xlabel('飞行方向 (度)')
# ax1.set_ylabel('飞行速度 (m/s)')
# ax1.set_title('飞行方向与速度对遮蔽时间的影响')
# plt.colorbar(sc1, ax=ax1, label='遮蔽时间 (s)')
# ax1.grid(True, alpha=0.3)

# 标记最佳点
# ax1.scatter(np.degrees(alpha_opt), v_opt, c='red', s=100, marker='*', label='最佳点')
# ax1.legend()

# 2. 2D散点图 - 投放时间 vs 起爆时间
# ax2 = fig.add_subplot(222)
# sc2 = ax2.scatter(samples[:, 2], samples[:, 3],
#                   c=shielding_times, cmap='plasma', s=50, alpha=0.8)
# ax2.set_xlabel('投放时间 (s)')
# ax2.set_ylabel('起爆时间 (s)')
# ax2.set_title('投放与起爆时间对遮蔽时间的影响')
# plt.colorbar(sc2, ax=ax2, label='遮蔽时间 (s)')
# ax2.grid(True, alpha=0.3)
#
# # 标记最佳点
# ax2.scatter(t_drop_opt, t_burst_opt, c='red', s=100, marker='*', label='最佳点')
# ax2.legend()
#
# # 3. 3D散点图 - 飞行方向、速度、遮蔽时间
# ax3 = fig.add_subplot(223, projection='3d')
# sc3 = ax3.scatter(np.degrees(samples[:, 0]), samples[:, 1], shielding_times,
#                   c=shielding_times, cmap='viridis', s=50, alpha=0.8)
# ax3.set_xlabel('飞行方向 (度)')
# ax3.set_ylabel('飞行速度 (m/s)')
# ax3.set_zlabel('遮蔽时间 (s)')
# ax3.set_title('飞行参数与遮蔽时间的关系')

# 4. 遮蔽时间分布直方图
ax4 = fig.add_subplot(224)
ax4.hist(shielding_times, bins=20, alpha=0.7, color='skyblue', edgecolor='black')
ax4.axvline(precise_shielding_time, color='red', linestyle='--',
            label=f'最佳遮蔽时间: {precise_shielding_time:.2f}s')
ax4.set_xlabel('遮蔽时间 (s)')
ax4.set_ylabel('频次')
ax4.set_title('遮蔽时间分布')
ax4.legend()
ax4.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# 创建参数重要性分析图
fig2, axes = plt.subplots(2, 2, figsize=(15, 12))

# 参数与遮蔽时间的相关性分析
param_names = ['飞行方向', '飞行速度', '投放时间', '起爆时间']
param_units = ['(度)', '(m/s)', '(s)', '(s)']

for i, ax in enumerate(axes.flatten()):
    if i < 4:
        # 计算相关性
        correlation = np.corrcoef(samples[:, i], shielding_times)[0, 1]

        ax.scatter(samples[:, i] if i != 0 else np.degrees(samples[:, 0]),
                   shielding_times, alpha=0.6)
        ax.set_xlabel(f'{param_names[i]} {param_units[i]}')
        ax.set_ylabel('遮蔽时间 (s)')
        ax.set_title(f'{param_names[i]} vs 遮蔽时间\n相关系数: {correlation:.3f}')
        ax.grid(True, alpha=0.3)

        # 添加趋势线
        if i == 0:
            x_vals = np.degrees(samples[:, 0])
        else:
            x_vals = samples[:, i]

        z = np.polyfit(x_vals, shielding_times, 1)
        p = np.poly1d(z)
        ax.plot(x_vals, p(x_vals), "r--", alpha=0.8)

plt.tight_layout()
plt.show()

# 输出统计信息
print("\n=== 统计分析 ===")
print(f"最佳遮蔽时间: {precise_shielding_time:.2f} s")
print(f"平均遮蔽时间: {np.mean(shielding_times):.2f} s")
print(f"最大遮蔽时间: {np.max(shielding_times):.2f} s")
print(f"最小遮蔽时间: {np.min(shielding_times):.2f} s")
print(f"标准差: {np.std(shielding_times):.2f} s")

print("\n=== 参数相关性 ===")
for i, name in enumerate(param_names):
    if i == 0:
        corr = np.corrcoef(np.degrees(samples[:, 0]), shielding_times)[0, 1]
    else:
        corr = np.corrcoef(samples[:, i], shielding_times)[0, 1]
    print(f"{name}: {corr:.3f}")

# 寻找次优解
top_5_indices = np.argsort(shielding_times)[-5:]
print("\n=== 前5个最佳解 ===")
for j, idx in enumerate(top_5_indices[::-1]):
    print(f"第{j + 1}名: 遮蔽时间 = {shielding_times[idx]:.2f}s")
    print(f"  飞行方向: {np.degrees(samples[idx, 0]):.2f}°")
    print(f"  飞行速度: {samples[idx, 1]:.2f} m/s")
    print(f"  投放时间: {samples[idx, 2]:.2f} s")
    print(f"  起爆时间: {samples[idx, 3]:.2f} s")
    print()

# 创建等高线图（飞行方向 vs 飞行速度）
if len(samples) > 0:
    fig3, ax = plt.subplots(figsize=(10, 8))

    # 创建网格数据进行插值
    xi = np.linspace(np.degrees(samples[:, 0].min()), np.degrees(samples[:, 0].max()), 100)
    yi = np.linspace(samples[:, 1].min(), samples[:, 1].max(), 100)
    Xi, Yi = np.meshgrid(xi, yi)

    # 插值计算遮蔽时间
    Zi = griddata((np.degrees(samples[:, 0]), samples[:, 1]), shielding_times, (Xi, Yi), method='cubic')

    # 绘制等高线
    contour = ax.contourf(Xi, Yi, Zi, levels=20, cmap='viridis', alpha=0.8)
    ax.contour(Xi, Yi, Zi, levels=10, colors='black', alpha=0.3, linewidths=0.5)

    # 标记最佳点
    ax.scatter(np.degrees(alpha_opt), v_opt, c='red', s=100, marker='*', label='最佳点')

    ax.set_xlabel('飞行方向 (度)')
    ax.set_ylabel('飞行速度 (m/s)')
    ax.set_title('飞行方向与速度对遮蔽时间的等高线图')
    plt.colorbar(contour, ax=ax, label='遮蔽时间 (s)')
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()